package com.zcq.aifitnessagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件持久化的 对话记忆实现类
 */
public class FileBasedChatMemory implements ChatMemory {

    private static final Kryo kryo = new Kryo();

    private final String BASE_DIR;

    static {
        kryo.setRegistrationRequired(false);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBasedChatMemory(String dir) {
        BASE_DIR = dir;
        File dirFile = new File(dir);
        if (!dirFile.exists()) {
            dirFile.mkdirs();
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        File file = getConversationFile(conversationId);
        List<Message> messagesList = new ArrayList<>();

        // 如果文件已存在，获取旧数据
        if (file.exists()) {
            messagesList.addAll(readConversation(file));
        }
        // 累加新输入
        messagesList.addAll(messages);
        // 写入
        saveConversation(file, messagesList);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messages = readConversation(getConversationFile(conversationId));
        return messages.stream().skip(Math.max(0, messages.size() - lastN)).toList();
    }

    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }

    /**
     * 获取会话对应的文件
     *
     * @param conversationId 会话id
     * @return File对象
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }

    /**
     * 读取 会话内容集合
     *
     * @param file 文件
     * @return 会话内容集合
     */
    private static List<Message> readConversation(File file) {
        if (file == null || !file.exists()) {
            return new ArrayList<>();
        }
        List<Message> messages;

        try (FileInputStream fileInputStream = new FileInputStream(file);
             Input input = new Input(fileInputStream);
        ) {
            messages = kryo.readObject(input, ArrayList.class);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        if (messages == null) {
            messages = new ArrayList<>();
        }

        return messages;
    }

    /**
     * 写入 会话内容
     *
     * @param file     文件
     * @param messages 会话内容
     */
    private static void saveConversation(File file, List<Message> messages) {
        try (FileOutputStream fileOutputStream = new FileOutputStream(file);
             Output output = new Output(fileOutputStream)
        ) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

}
